import os
import random
import matplotlib as mpl
import matplotlib.pyplot as plt
import numpy as np


# 可视化
class Tranning_Visualization:
    epochs_losses = []
    epoch_val_loss = []
    # {epoch:{'map':0.0, 'map_50':0.0, 'map_75':0.0, 'map_small':0.0, 'map_medium':0.0, 'map_large':0.0}}
    epochs_mAP = {}           
    mAP_keys = ['map', 'map_50', 'map_75', 'map_small', 'map_medium', 'map_large']
    def __init__(self, save_dir='resault'):
        self.save_dir = save_dir

        self.mAP_fig_name = ''
        self.losses_fig_name = ''

    def draw_epochs_losses(self, name='epochs_losses', cmap_name='Set3'):
        self.losses_fig_name = name
        self.losses_fig = plt.figure(self.losses_fig_name)
        plt.title('epoch-loss', loc='center', pad=20)
        color_num = 2
        
        # 调色色组
        cmap = mpl.cm.get_cmap(cmap_name, color_num)
        colors_list = cmap(np.linspace(0, 1, color_num))
        plt.xlabel('epoch')
        plt.ylabel('loss')
        plt.plot(np.arange(0, len(self.epochs_losses), 1), np.array(self.epochs_losses), color=colors_list[0, :], label='train loss')
        plt.plot(np.arange(0, len(self.epoch_val_loss), 1), np.array(self.epoch_val_loss), color=colors_list[1, :], label='val loss')
        plt.legend(loc='upper right')

    def draw_epochs_mAP(self, name='epochs_mAP', cmap_name='Set3', if_merge=True):
        self.mAP_fig_name = name
        self.mAP_fig = plt.figure(self.mAP_fig_name)
        plt.title('epoch-mAP', loc='center')
        if if_merge:
            self.draw_mAP_merge(name, cmap_name)
        else:
            self.draw_mAP_subplot(name, cmap_name)

    def draw_mAP_merge(self, name, cmap_name='Set3'):
        nrow = 2
        ncol = 3
        color_num = nrow * ncol
        cmap = mpl.cm.get_cmap(cmap_name, color_num)
        colors_list = cmap(np.linspace(0, 1, color_num))
        if self.epochs_mAP:
            for i, key in enumerate(self.mAP_keys):
                if i > nrow * ncol:
                    break
                mAPs = []
                epochs = []
                for e in self.epochs_mAP.keys():
                    mAPs.append(float(self.epochs_mAP[e][key]))
                    epochs.append(e)
                plt.plot(np.array(epochs), np.array(mAPs), color=colors_list[i, :], label=key)
                plt.xlabel('epoch')
                plt.ylabel('map')
            plt.legend(loc='lower right')

    def draw_mAP_subplot(self, name, cmap_name='Set3'):
        nrow = 2
        ncol = 3
        color_num = 1
        cmap = mpl.cm.get_cmap(cmap_name, color_num)
        colors_list = cmap(np.linspace(0, 1, color_num))
        if self.epochs_mAP:
            # _d = list(self.epochs_mAP.items())[0]
            for i, key in enumerate(self.mAP_keys):
                if i > nrow * ncol:
                    break
                plt.subplot(nrow, ncol, i)
                mAPs = []
                epochs = []
                for e in self.epochs_mAP.keys():
                    mAPs.append(float(self.epochs_mAP[e][key]))
                    epochs.append(e)
                plt.plot(np.array(epochs), np.array(mAPs), color=colors_list[0, :], label=key)
                plt.xlabel('epoch')
                plt.ylabel('map')

    def save_all(self):
        if self.losses_fig_name:
            plt.figure(self.losses_fig_name)
            plt.savefig(os.path.join(self.save_dir, 'epochs_losses.png'))

        if self.mAP_fig_name:
            plt.figure(self.mAP_fig_name)
            plt.savefig(os.path.join(self.save_dir, 'epochs_mAP.png'))


if __name__ == '__main__':
    tv = Tranning_Visualization()
    train_losse = np.random.randn(100)
    val_losse = np.random.randn(100)
    tv.epochs_losses = train_losse.tolist()
    tv.epoch_val_loss = val_losse.tolist()
    tv.draw_epochs_losses()
    tv.save_all()